import numpy as np
from sklearn.metrics import precision_score, make_scorer
from sklearn.base import BaseEstimator


# 自定义评估器类
class WeightedFusion(BaseEstimator):
    def __init__(self, w_cos=0.7, w_bm25=0.3, threshold=0.7):
        self.w_cos = w_cos
        self.w_bm25 = w_bm25
        self.threshold = threshold

    def fit(self, X, y=None):
        # 不需要实际训练，仅用于参数搜索
        return self

    def score(self, X, y):
        scores = X[:, 0] * self.w_cos + X[:, 1] * self.w_bm25
        pred = (scores > self.threshold).astype(int)
        return precision_score(y, pred)  # 直接最大化精确率


# 模拟验证集数据
np.random.seed(42)
n_samples = 2000
s_cos = np.clip(np.random.normal(0.7, 0.15, n_samples), 0, 1)
s_bm25 = np.clip(np.random.beta(2, 3, n_samples), 0, 1)
y_true = np.where((s_cos > 0.65) & (s_bm25 > 0.4), 1, 0)
X_val = np.column_stack([s_cos, s_bm25])

# 创建参数搜索器
model = WeightedFusion()
param_grid = {
    'w_cos': np.linspace(0.5, 0.8, 15),
    'w_bm25': np.linspace(0.2, 0.5, 15)
}

# 使用自定义评分
precision_scorer = make_scorer(precision_score)

# 执行网格搜索
from sklearn.model_selection import GridSearchCV

grid = GridSearchCV(estimator=model,
                    param_grid=param_grid,
                    scoring=precision_scorer,
                    cv=5,
                    n_jobs=-1)

grid.fit(X_val, y_true)

# 输出最佳参数
print(f"Best Weights: w_cos={grid.best_params_['w_cos']:.3f}, w_bm25={grid.best_params_['w_bm25']:.3f}")
print(f"Best Precision: {grid.best_score_:.2%}")

# 验证最终性能